import loss, nn
import numpy as np
import utils
import matplotlib.pyplot as plt

"""
Highest API Model of network
"""


class Model:
    def __init__(self, network, auto_grad=False):
        self.network = network
        self.optimizer = None
        self.history_acc = []
        self.epochs = None
        self.test_acc = None

    def __log_division(self, epoch, step, loss, accuracy):
        print("epoch:"+str(epoch)+", step:"+ str(step)+", loss_value:"+str(loss)+", accuracy:"+str(accuracy))

    def __get_accuracy_division(self, x, y):
        sum_V = 0
        for i in range(len(x)):
            outputs, loss_v = self.network(x[i], y[i])
            if outputs.argmax() == y[i].argmax():
                sum_V += 1
        return sum_V / len(x)

    def compile(self, optimizer):
        self.optimizer = optimizer

    def train(self, x, y, epochs, batch_size=32):
        self.epochs = epochs
        if len(x) != len(y):
            raise ValueError("The length of x and y should be equal.")
        for i in range(epochs):
            accu = 0
            counter = 0
            while counter < len(x):
                temp = np.random.random_integers(0, len(x) - 1, batch_size)
                for j in temp:
                    counter += 1
                    outputs, loss_value = self.network(x[j], y[j])
                    self.network.backward()
                    self.network.updateParams(optimizer=self.optimizer.step)
                    if outputs.argmax() == y[j].argmax():
                        accu += 1 /len(x)
                    print(end=f'\repochs {i}, progress: {round(counter/(len(x))*100,5)}%, train_accu: {round(accu*100, 5)}%')
            self.history_acc.append(accu)
            print(end='\n')

    def fit(self, x, y):
        self.test_acc = self.__get_accuracy_division(x, y)
        print("test_accu:", self.test_acc)

    def drawHistory(self):
        x = np.arange(self.epochs)
        plt.plot(x, self.history_acc, marker='o', label='train', markevery=2)
        plt.plot(x, np.full(self.epochs, self.test_acc), marker='s', label='test', markevery=2)
        plt.xlabel("epochs")
        plt.ylabel("train_accuracy")
        plt.ylim(0, 1.0)
        plt.legend(loc='lower right')
        plt.show()






